import cv2
import time
import numpy as np
from stereo.utils.stereoRectify_process import stereo_rectify
from stereo.utils.config import Config
import threading
import queue
import sys
import os

class StereoCameraCaptureThreadedDual:
    def __init__(self, left_device, right_device, width=640, height=480, fps=30, max_queue_size=5):
        self.left_device = left_device
        self.right_device = right_device
        self.width = width
        self.height = height
        self.fps = fps
        self.max_queue_size = max_queue_size
        self.left_queue = queue.Queue(maxsize=max_queue_size)
        self.right_queue = queue.Queue(maxsize=max_queue_size)
        self.frame_queue = queue.Queue(maxsize=max_queue_size)
        self.cap_left = self._open_camera(self.left_device)
        self.cap_right = self._open_camera(self.right_device)
        self.running = False
        self.left_thread = None
        self.right_thread = None
        self.sync_thread = None
        self.capture_fps = fps
        self.capture_interval = 1.0 / self.capture_fps if self.capture_fps > 0 else 0
        self.last_capture_time = 0
        # 加载相机标定参数
        self.map1x, self.map1y, self.map2x, self.map2y = stereo_rectify.map1x, stereo_rectify.map1y, stereo_rectify.map2x, stereo_rectify.map2y
        self.start_capture()

    def _open_camera(self, device_name):
        gst_pipeline = (
            f'v4l2src device=/dev/{device_name} ! '
            f'video/x-raw,format=NV12,width={self.width},height={self.height},framerate={self.fps}/1 ! '
            f'appsink'
        )
        cap = cv2.VideoCapture(gst_pipeline, cv2.CAP_GSTREAMER)
        if not cap.isOpened():
            print(f"无法打开摄像头: {device_name}")
            return None
        actual_width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
        actual_height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
        print(f"[{device_name}] 实际帧大小: {actual_width}x{actual_height}")
        return cap

    def apply_rectification_single(self, img, map1x, map1y, k=None, dist=None):
        if k is not None:
            img = cv2.undistort(img, k, dist)
        rectified = cv2.remap(img, map1x, map1y, cv2.INTER_LINEAR)
        return rectified

    def _process_frame(self, frame, is_left=True):
        try:
            frame = cv2.cvtColor(frame, cv2.COLOR_YUV2BGR_NV12)
            frame = cv2.rotate(frame, cv2.ROTATE_180)
            frame = self.apply_rectification_single(
                frame,
                self.map1x if is_left else self.map2x,
                self.map1y if is_left else self.map2y,
            )
            return frame
        except Exception as e:
            print(f"图像处理失败: {e}")
            return None

    def _left_capture_loop(self):
        while self.running:
            if self.left_queue.full():
                time.sleep(0.005)
                continue
            ret, frame = self.cap_left.read()
            timestamp = time.time()
            if ret:
                processed = self._process_frame(frame, True)
                if processed is not None:
                    try:
                        self.left_queue.put((timestamp, processed), block=False)
                    except queue.Full:
                        pass
            else:
                print("未能从左摄像头读取帧")
                time.sleep(0.01)

    def _right_capture_loop(self):
        while self.running:
            if self.right_queue.full():
                time.sleep(0.005)
                continue
            ret, frame = self.cap_right.read()
            timestamp = time.time()
            if ret:
                processed = self._process_frame(frame, False)
                if processed is not None:
                    try:
                        self.right_queue.put((timestamp, processed), block=False)
                    except queue.Full:
                        pass
            else:
                print("未能从右摄像头读取帧")
                time.sleep(0.01)

    def _sync_loop(self):
        # 简单同步：取出时间戳最接近的左右帧
        while self.running:
            try:
                left_ts, left_frame = self.left_queue.get(timeout=0.1)
                right_ts, right_frame = self.right_queue.get(timeout=0.1)
                # 简单同步策略：只要时间差小于阈值就配对
                if abs(left_ts - right_ts) < 0.03:  # 30ms内认为同步
                    try:
                        self.frame_queue.put((left_frame, right_frame, (left_ts + right_ts) / 2), block=False)
                    except queue.Full:
                        pass
                else:
                    # 时间差太大，丢弃较早帧
                    if left_ts < right_ts:
                        # 丢弃左帧，右帧重新入队
                        self.right_queue.put((right_ts, right_frame))
                    else:
                        self.left_queue.put((left_ts, left_frame))
            except queue.Empty:
                continue

    def start_capture(self):
        self.running = True
        self.left_thread = threading.Thread(target=self._left_capture_loop, daemon=True)
        self.right_thread = threading.Thread(target=self._right_capture_loop, daemon=True)
        self.sync_thread = threading.Thread(target=self._sync_loop, daemon=True)
        self.left_thread.start()
        self.right_thread.start()
        self.sync_thread.start()

    def set_capture_fps(self, capture_fps):
        if capture_fps <= 0:
            print("警告: 捕获帧率必须大于0，设置为默认值1")
            capture_fps = 1
        self.capture_fps = capture_fps
        self.capture_interval = 1.0 / self.capture_fps
        print(f"捕获帧率已设置为: {self.capture_fps} fps (帧间隔: {self.capture_interval*1000:.2f}ms)")
        return True

    def get_stereo_frames(self, block=True, timeout=1.0):
        try:
            if block:
                frame_data = self.frame_queue.get(timeout=timeout)
            else:
                if self.frame_queue.empty():
                    return None, None, None
                frame_data = self.frame_queue.get_nowait()
            return frame_data[0], frame_data[1], frame_data[2]
        except queue.Empty:
            print("获取帧超时")
            return None, None, None

    def get_queue_size(self):
        return self.frame_queue.qsize()

    def clear_queue(self):
        while not self.frame_queue.empty():
            try:
                self.frame_queue.get_nowait()
            except queue.Empty:
                break
        return True

    def close(self):
        self.running = False
        if self.left_thread:
            self.left_thread.join(timeout=1.0)
        if self.right_thread:
            self.right_thread.join(timeout=1.0)
        if self.sync_thread:
            self.sync_thread.join(timeout=1.0)
        if self.cap_left:
            self.cap_left.release()
        if self.cap_right:
            self.cap_right.release()
        print("摄像头已释放")

    def __del__(self):
        self.close()

StereoCameraCapture_threaded_dual_ins = StereoCameraCaptureThreadedDual("video22", "video31")
StereoCameraCapture_threaded_dual_ins.set_capture_fps(25)

if __name__ == "__main__":
    save_video = False
    target_fps = 30
    StereoCameraCapture_threaded_dual_ins.set_capture_fps(target_fps)
    if save_video:
        fourcc = cv2.VideoWriter_fourcc(*'XVID')
        out = cv2.VideoWriter('stereo_output_dual.avi', fourcc, target_fps, (640, 480*2))
        try:
            frame_count = 0
            start_time = time.time()
            while True:
                frame_left, frame_right, frame_time = StereoCameraCapture_threaded_dual_ins.get_stereo_frames()
                if frame_left is None or frame_right is None:
                    print("未能获取有效帧")
                    continue
                combined_frame = np.vstack((frame_left, frame_right))
                out.write(combined_frame)
                frame_count += 1
                elapsed = time.time() - start_time
                if elapsed >= 5.0:
                    actual_fps = frame_count / elapsed
                    print(f"实际帧率: {actual_fps:.2f} fps (目标: {target_fps} fps)")
                    print(f"队列大小: {StereoCameraCapture_threaded_dual_ins.get_queue_size()}")
                    frame_count = 0
                    start_time = time.time()
        except KeyboardInterrupt:
            print("程序被用户中断")
        finally:
            out.release()
            cv2.destroyAllWindows()
            StereoCameraCapture_threaded_dual_ins.close()
    else:
        try:
            frame_count = 0
            start_time = time.time()
            while True:
                frame_left, frame_right, frame_time = StereoCameraCapture_threaded_dual_ins.get_stereo_frames()
                if frame_left is None or frame_right is None:
                    continue
                frame_count += 1
                elapsed = time.time() - start_time
        except KeyboardInterrupt:
            print("程序被用户中断")
        finally:
            StereoCameraCapture_threaded_dual_ins.close() 